import akshare as ak
import pandas as pd
import sqlite3
import os

DB_FILE = "stock_data.db"
FINANCIAL_DB_FILE = "financial_data.db"

def get_stock_data(code: str, start_date: str, end_date: str) -> pd.DataFrame:
    """
    获取股票历史数据，优先从本地SQLite数据库读取。
    如果缓存无效或不完整，则通过akshare重新下载并替换缓存。
    """
    table_name = f"stock_{code.replace('.', '_')}"
    db_path = os.path.join(os.path.dirname(__file__), '..', '..', DB_FILE)
    
    req_start = pd.to_datetime(start_date).strftime('%Y-%m-%d')
    req_end = pd.to_datetime(end_date).strftime('%Y-%m-%d')

    try:
        with sqlite3.connect(db_path) as conn:
            # 检查表是否存在
            cursor = conn.cursor()
            cursor.execute("SELECT name FROM sqlite_master WHERE type='table' AND name=?", (table_name,))
            if cursor.fetchone():
                query = f"SELECT min(日期) as min_date, max(日期) as max_date FROM {table_name}"
                cached_range = pd.read_sql(query, conn).iloc[0]
                if cached_range['min_date'] <= req_start and cached_range['max_date'] >= req_end:
                    print(f"从数据库缓存加载 {code} 数据。")
                    final_query = f"SELECT * FROM {table_name} WHERE 日期 BETWEEN '{req_start}' AND '{req_end}'"
                    return pd.read_sql(final_query, conn)
    except Exception as e:
        print(f"检查数据库缓存时出错，将重新下载: {e}")

    print(f"缓存无效或不完整，从 akshare 下载 {code} 数据...")
    try:
        df_live = ak.stock_zh_a_hist(symbol=code, period="daily", start_date=start_date.replace("-", ""), end_date=end_date.replace("-", ""), adjust="qfq")
        if df_live.empty:
            print(f"akshare 未能获取到 {code} 的数据。")
            return pd.DataFrame()
            
        with sqlite3.connect(db_path) as conn:
            # 使用 'replace' 模式，确保数据干净
            df_live.to_sql(table_name, conn, if_exists='replace', index=False)
            print(f"数据已下载并缓存到 {table_name} 表。")
        
        # 再次查询以确保返回正确的时间范围
        final_query = f"SELECT * FROM {table_name} WHERE 日期 BETWEEN '{req_start}' AND '{req_end}'"
        with sqlite3.connect(db_path) as conn:
             return pd.read_sql(final_query, conn)

    except Exception as e:
        print(f"从 akshare 下载数据时发生严重错误: {e}")
        return pd.DataFrame() 

def get_financial_data(code: str) -> pd.DataFrame:
    """
    获取并缓存公司的主要财务指标数据。
    """
    table_name = f"financial_{code.replace('.', '_')}"
    
    # 优先从缓存读取
    if os.path.exists(FINANCIAL_DB_FILE):
        try:
            with sqlite3.connect(FINANCIAL_DB_FILE) as conn:
                df = pd.read_sql(f"SELECT * FROM {table_name}", conn)
                if not df.empty:
                    print(f"从缓存加载 {code} 的财务数据。")
                    return df
        except Exception as e:
            print(f"从缓存读取财务数据失败: {e}")
            
    # 缓存不存在或读取失败，则从 akshare 下载
    print(f"从 akshare 下载 {code} 的财务数据...")
    try:
        df_financial = ak.stock_financial_analysis_indicator(symbol=code)
        if df_financial.empty:
            return pd.DataFrame()
        
        with sqlite3.connect(FINANCIAL_DB_FILE) as conn:
            df_financial.to_sql(table_name, conn, if_exists='replace', index=False)
        
        return df_financial
    except Exception as e:
        print(f"从 akshare 下载财务数据失败: {e}")
        return pd.DataFrame() 